{-# LANGUAGE ForeignFunctionInterface #-}
module Hexena.Cmds.GetSVMData (run) where

import Data.Binary.Get
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Data.IFF
import qualified Data.Map as M
import Data.Maybe (fromJust)
--import System.Environment (getArgs)
import Text.Printf
import Data.List.Split (splitOn)

import Hexena.Chunks
import Hexena.MZHeader
import Hexena.PEHeader
import Hexena.Math.Statistics (calculateE)

lfanew :: B.ByteString -> Double
lfanew mz = if lf < 0x800 then fromIntegral lf / 0x800 else 1.0
    where lf = runGet getWord32be $ BL.fromChunks [B.drop 60 mz]

numberOfSections :: PEHeader -> Double
numberOfSections pe = fromIntegral (min sects 32) / 32.0
    where sects = peNumberOfSections pe

-- | Create a single line for the libsvm-tools.
svmtransformation :: IFFMap -> [T] -> String
svmtransformation iff _ = unwords $ zipWith (printf "%d:%.8e") [(1::Int)..] vals
    where mz = content $ fromJust $ M.lookup "MZHD" iff
          mz' = content $ fromJust $ M.lookup "MZRW" iff
          mzDist = mzDistanceN $ B.unpack mz'
          (hist, histC) = loadFORMHIST iff
          entr1 = calculateE hist
          entr1C = calculateE histC
          vals = [ numberOfSections $ loadPEHD iff
                 , mzDist
                 , lfanew mz
                 , entr1 / 8.0
                 , entr1C / 8.0
                 ]
                 -- The following did not produce a good model...
                 -- ++ elems (histToProbability hist)
                 -- ++ elems (histToProbability histC)

getDefaultClass :: String -> Int
getDefaultClass = read

-- One-class SVM generation from hexena files.
run :: Int -> [FilePath] -> IO ()
run classifier filenames = do
  -- get the default classifier for SVM as a string
  let dc = show classifier ++ "  "
  dats <- mapM loadMalwIFF' filenames
  putStr $ unlines $ map ((++) dc . uncurry svmtransformation) dats
